package main

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
	"encoding/csv"
	"context"
	"os"
	"sync"

	"golang.org/x/net/dns/dnsmessage"
)

var namedGroupsToName = map[uint16]string{
	uint16(tls.HybridSIDHp503Curve25519): "X25519-SIDHp503",
	uint16(tls.HybridSIKEp503Curve25519): "X25519-SIKEp503",
	uint16(tls.X25519):                   "X25519",
	uint16(tls.CurveP256):                "P-256",
	uint16(tls.CurveP384):                "P-384",
	uint16(tls.CurveP521):                "P-521",
}

var cipherSuiteIdToName = map[uint16]string{
	tls.TLS_RSA_WITH_AES_128_CBC_SHA:            "TLS_RSA_WITH_AES_128_CBC_SHA",
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256: "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
	tls.TLS_AES_128_GCM_SHA256:                  "TLS_AES_128_GCM_SHA256",
	tls.TLS_AES_256_GCM_SHA384:                  "TLS_AES_256_GCM_SHA384",
	tls.TLS_CHACHA20_POLY1305_SHA256:            "TLS_CHACHA20_POLY1305_SHA256",
}

func main(){
	if len(os.Args) < 2 {
		log.Fatal("Usage: go run main.go <csv_file>")
	}
	csvFile := os.Args[1]

	// 打开 CSV 文件
	file, err := os.Open(csvFile)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	// 创建CSV阅读器
	reader := csv.NewReader(file)
	var domains []string

	// 读取 CSV 文件中的所有域名
	for {
		record, err := reader.Read()
		if err != nil {
			break
		}
		domains = append(domains, record[0])
	}

	// 使用 WaitGroup 控制并发，每次最多运行 3 个任务
	var wg sync.WaitGroup
	sem := make(chan struct{}, 3)

	for _, domain := range domains {
		wg.Add(1)
		sem <- struct{}{} // 限制并发数

		go func(d string) {
			defer wg.Done()
			esniClientStart("www." + d)
			<-sem // 释放并发限制
			time.Sleep(500 * time.Millisecond)
		}(domain)
	}

	wg.Wait()

}

func esniClientStart(domain string) {

	//读取并解析ESNI公钥
	contents, err := ioutil.ReadFile("esni.pub")
	if err != nil {
		log.Fatalf("Failed to read ESNIKeys: %s", err)
	}
	esniKeysBytes, err := base64.StdEncoding.DecodeString(string(contents))
	if err != nil {
		log.Fatalf("Failed to parse -esni-keys: %s", err)
	}
	clientESNIKeys,err:=tls.ParseESNIKeys(esniKeysBytes)
	if clientESNIKeys == nil {
		log.Fatalf("Failed to process ESNI response for host: %s", err)
	}

	//配置客户端证书
	caCert, err := ioutil.ReadFile("server1.crt")
	if err != nil {
		fmt.Println(err)
		return
	}
	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	//配置TLS
	tlsConfig := &tls.Config{
		MinVersion: tls.VersionTLS13,
		MaxVersion: tls.VersionTLS13,
		ServerName: "server1.com",
		ClientESNIKeys:clientESNIKeys,
		RootCAs:caCertPool,
	}

	//创建http.client实例
	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: tlsConfig,
			DisableKeepAlives: true,
		},
	}

	// 构建 DNS 请求
	dnsReq := buildDNSRequest(domain)

	// 创建一个5秒的超时上下文
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// 使用超时上下文发送 ESNI 请求
	resp, err := sendDoHRequest(client, ctx, "https://192.168.157.128/dns-query", dnsReq)
	if err != nil {
		// 忽略超时错误，不终止程序
		if strings.Contains(err.Error(), "context deadline exceeded") {
			log.Printf("Timeout error: ESNI request to %s for domain %s exceeded time limit", "https://192.168.157.128/dns-query", domain)
			return
		}
		log.Printf("Error sending ESNI request: %v", err)
		return
	}
	if resp == nil {
		log.Printf("Received nil response for domain: %s", domain)
		return
	}
	defer resp.Body.Close()

	// 解析响应
	err = parseDNSResponse(resp.Body)
	if err != nil {
		log.Printf("Error parsing DNS response for domain %s: %v", domain, err)
	}

}

//根据别名获取实值
func getIDByName(m map[uint16]string, name string) (uint16, error) {
	for key, value := range m {
		if value == name {
			return key, nil
		}
	}
	return 0, errors.New("Unknown value")
}

// 构建 DNS 请求，查询指定域名
func buildDNSRequest(domain string) []byte {
	// 确保域名以 . 结尾
	if !strings.HasSuffix(domain, ".") {
		domain += "."
	}

	// 构建 DNS 消息
	msg := dnsmessage.Message{
		Header: dnsmessage.Header{
			RecursionDesired: true,
		},
		Questions: []dnsmessage.Question{
			{
				Name:  dnsmessage.MustNewName(domain),
				Type:  dnsmessage.TypeA,
				Class: dnsmessage.ClassINET,
			},
		},
	}

	// 序列化 DNS 消息
	dnsBytes, err := msg.Pack()
	if err != nil {
		log.Fatal("Failed to pack DNS message: ", err)
	}
	return dnsBytes
}

// 发送 DoH 请求
func sendDoHRequest(client *http.Client, ctx context.Context, url string, dnsRequest []byte) (*http.Response, error) {
	// 创建 HTTP POST 请求，正文包含 DNS 请求的二进制数据
	req, err := http.NewRequest("POST", url, bytes.NewReader(dnsRequest))
	if err != nil {
		return nil, err
	}

	// 设置 HTTP 头部，指定内容类型为 DNS 消息格式
	req.Header.Set("Content-Type", "application/dns-message")
	req = req.WithContext(ctx)

	// 发送请求
	return client.Do(req)
}

// 解析 DNS 响应
func parseDNSResponse(responseBody io.Reader) error {
	// 读取响应的二进制数据
	dnsRespBytes, err := ioutil.ReadAll(responseBody)
	if err != nil {
		return fmt.Errorf("failed to read response body: %v", err)
	}

	// 解析 DNS 响应
	var msg dnsmessage.Message
	err = msg.Unpack(dnsRespBytes)
	if err != nil {
		return fmt.Errorf("failed to unpack DNS response: %v", err)
	}

	// 打印响应中的查询结果
	for _, answer := range msg.Answers {
		switch answer.Body.(type) {
		case *dnsmessage.AResource:
			a := answer.Body.(*dnsmessage.AResource)
			fmt.Printf("Domain: %s -> IP Address: %v\n", msg.Questions[0].Name, a.A)
		default:
			fmt.Printf("Domain: %s -> Other record: %v\n", msg.Questions[0].Name, answer.Body)
		}
	}

	return nil
}
